#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
视频下载器模块
提供从微信文章URL中提取和下载视频的功能
"""

import os
import re
import time
import logging
import requests
import traceback
import pymysql
from pathlib import Path
from datetime import datetime
from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeoutError
from urllib.parse import urlparse, urljoin
import random

# 导入配置
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import DB_CONFIG, VIDEO_CONFIG, BROWSER_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class VideoDownloader:
    """视频下载器类"""
    
    def __init__(self):
        """初始化视频下载器"""
        self.download_dir = VIDEO_CONFIG['download_dir']
        self.max_retry = VIDEO_CONFIG.get('max_retry', 3)
        self.timeout = VIDEO_CONFIG.get('timeout', 60)
        
        # 确保下载目录存在
        self.download_dir.mkdir(parents=True, exist_ok=True)
        
        logger.info(f"视频下载器初始化完成，下载目录: {self.download_dir}")
    
    def download_videos_from_article(self, article_url, article_title=None):
        """
        从文章URL中下载所有视频
        
        Args:
            article_url (str): 文章URL
            article_title (str, optional): 文章标题，如果不提供会自动提取
            
        Returns:
            dict: 下载结果
        """
        logger.info(f"开始处理文章: {article_url}")
        
        result = {
            'success': True,
            'article_url': article_url,
            'article_title': article_title,
            'video_urls': [],
            'downloaded_videos': [],
            'failed_videos': [],
            'error': None
        }
        
        try:
            # 如果没有提供标题，尝试从URL或页面提取
            if not article_title:
                article_title = self._extract_title_from_url(article_url)
                result['article_title'] = article_title
            
            # 提取视频URL
            video_urls = self._extract_video_urls(article_url, article_title)
            result['video_urls'] = video_urls
            
            if not video_urls:
                logger.warning("未找到任何视频URL")
                return result
            
            logger.info(f"找到 {len(video_urls)} 个视频，开始下载...")
            
            # 下载每个视频
            for i, video_url in enumerate(video_urls):
                try:
                    download_result = self._download_single_video(
                        video_url, article_title, i
                    )
                    
                    if download_result['success']:
                        result['downloaded_videos'].append(download_result)
                        logger.info(f"视频 {i+1}/{len(video_urls)} 下载成功")
                    else:
                        result['failed_videos'].append({
                            'url': video_url,
                            'error': download_result['error']
                        })
                        logger.error(f"视频 {i+1}/{len(video_urls)} 下载失败: {download_result['error']}")
                    
                    # 添加延迟避免请求过快
                    time.sleep(random.uniform(1, 3))
                    
                except Exception as e:
                    error_msg = f"下载视频时发生错误: {str(e)}"
                    logger.error(error_msg)
                    result['failed_videos'].append({
                        'url': video_url,
                        'error': error_msg
                    })
            
            # 记录下载统计
            self._log_download_summary(result)
            
        except Exception as e:
            result['success'] = False
            result['error'] = str(e)
            logger.error(f"处理文章时发生错误: {e}")
            logger.error(traceback.format_exc())
        
        return result
    
    def _extract_video_urls(self, article_url, article_title):
        """使用Playwright提取文章中的视频URL"""
        video_urls = []
        
        try:
            logger.info("启动Playwright浏览器...")
            with sync_playwright() as p:
                # 启动浏览器
                browser = p.firefox.launch(
                    headless=BROWSER_CONFIG.get('headless', True),
                    slow_mo=BROWSER_CONFIG.get('slow_mo', 100)
                )
                logger.info("浏览器启动成功")
                
                page = browser.new_page()
                page.set_default_timeout(15000)  # 减少超时时间到15秒
                
                # 设置User-Agent
                page.set_extra_http_headers({
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                })
                
                logger.info(f"正在访问文章页面: {article_url}")
                try:
                    # 使用更短的超时时间和更简单的等待策略
                    page.goto(article_url, wait_until='domcontentloaded', timeout=15000)
                    logger.info("页面加载完成，等待内容渲染...")
                    
                    # 等待页面内容加载
                    time.sleep(2)
                    
                except Exception as e:
                    logger.error(f"页面访问失败: {e}")
                    browser.close()
                    return []
                
                # 方法1: 查找video标签
                try:
                    logger.info("方法1: 查找video标签...")
                    video_elements = page.query_selector_all('video')
                    logger.info(f"找到 {len(video_elements)} 个video标签")
                    
                    for video in video_elements:
                        src = video.get_attribute('src')
                        if src and src.startswith('http'):
                            video_urls.append(src)
                            logger.info(f"找到video标签视频: {src[:50]}...")
                except Exception as e:
                    logger.warning(f"提取video标签失败: {e}")
                
                # 方法2: 查找iframe中的视频
                try:
                    logger.info("方法2: 查找iframe视频...")
                    iframe_elements = page.query_selector_all('iframe')
                    logger.info(f"找到 {len(iframe_elements)} 个iframe")
                    
                    for iframe in iframe_elements:
                        src = iframe.get_attribute('src')
                        if src and ('video' in src.lower() or 'mp4' in src.lower()):
                            video_urls.append(src)
                            logger.info(f"找到iframe视频: {src[:50]}...")
                except Exception as e:
                    logger.warning(f"提取iframe视频失败: {e}")
                
                # 方法3: 从页面源码中使用正则表达式提取
                try:
                    logger.info("方法3: 正则表达式提取...")
                    page_content = page.content()
                    logger.info(f"页面内容长度: {len(page_content)} 字符")
                    
                    # 视频URL的正则表达式模式
                    video_patterns = [
                        r'https://[^"\']*\.mp4[^"\']*',
                        r'https://[^"\']*\.m3u8[^"\']*',
                        r'https://[^"\']*\.flv[^"\']*',
                        r'data-src="([^"]*\.mp4[^"]*)"',
                        r'src="([^"]*\.mp4[^"]*)"',
                        r'https://mmbiz\.qpic\.cn/[^"\']*\.mp4',
                    ]
                    
                    for pattern in video_patterns:
                        matches = re.findall(pattern, page_content, re.IGNORECASE)
                        logger.info(f"模式 {pattern} 找到 {len(matches)} 个匹配")
                        for match in matches:
                            video_url = match if isinstance(match, str) else match[0]
                            if self._is_valid_video_url(video_url):
                                video_urls.append(video_url)
                                logger.info(f"通过正则表达式找到视频: {video_url[:50]}...")
                
                except Exception as e:
                    logger.warning(f"正则表达式提取视频失败: {e}")
                
                # 方法4: 查找视频号卡片
                try:
                    logger.info("方法4: 查找视频号卡片...")
                    video_cards = page.query_selector_all('[data-url*="video"]')
                    logger.info(f"找到 {len(video_cards)} 个视频卡片")
                    
                    for card in video_cards:
                        video_url = card.get_attribute('data-url')
                        if video_url:
                            video_urls.append(video_url)
                            logger.info(f"找到视频号卡片: {video_url[:50]}...")
                except Exception as e:
                    logger.warning(f"提取视频号卡片失败: {e}")
                
                logger.info("关闭浏览器...")
                browser.close()
                
        except Exception as e:
            logger.error(f"提取视频URL时发生错误: {e}")
            logger.error(traceback.format_exc())
        
        # 去重并过滤
        unique_video_urls = list(set(video_urls))
        valid_video_urls = [url for url in unique_video_urls if self._is_valid_video_url(url)]
        
        logger.info(f"总共提取到 {len(valid_video_urls)} 个有效视频URL")
        return valid_video_urls
    
    def _is_valid_video_url(self, url):
        """检查是否是有效的视频URL"""
        if not url or len(url) < 10:
            return False
        
        # 检查是否是HTTP/HTTPS URL
        if not (url.startswith('http://') or url.startswith('https://')):
            return False
        
        # 检查文件扩展名或包含视频关键词
        video_indicators = ['.mp4', '.m3u8', '.flv', '.avi', '.mov', '.mkv', '.webm', 'video']
        url_lower = url.lower()
        
        return any(indicator in url_lower for indicator in video_indicators)
    
    def _download_single_video(self, video_url, article_title, index=0):
        """下载单个视频文件"""
        result = {
            'success': False,
            'video_url': video_url,
            'local_path': None,
            'file_size': 0,
            'error': None
        }
        
        try:
            # 创建安全的文件名
            safe_title = self._sanitize_filename(article_title)
            video_dir = self.download_dir / safe_title
            video_dir.mkdir(parents=True, exist_ok=True)
            
            # 生成文件名
            file_extension = self._get_video_extension(video_url)
            filename = f"video_{index+1:03d}{file_extension}"
            file_path = video_dir / filename
            
            # 如果文件已存在且大小合理，跳过下载
            if file_path.exists() and file_path.stat().st_size > 1024:
                logger.info(f"视频文件已存在，跳过下载: {filename}")
                result['success'] = True
                result['local_path'] = str(file_path)
                result['file_size'] = file_path.stat().st_size
                return result
            
            # 开始下载
            logger.info(f"开始下载视频: {video_url}")
            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Referer': 'https://mp.weixin.qq.com/',
                'Accept': 'video/webm,video/ogg,video/*;q=0.9,application/ogg;q=0.7,audio/*;q=0.6,*/*;q=0.5',
                'Accept-Encoding': 'identity',
                'Range': 'bytes=0-'
            }
            
            # 尝试下载，支持重试
            for attempt in range(self.max_retry):
                try:
                    response = requests.get(
                        video_url,
                        headers=headers,
                        stream=True,
                        timeout=self.timeout,
                        allow_redirects=True
                    )
                    response.raise_for_status()
                    
                    # 获取文件大小
                    total_size = int(response.headers.get('content-length', 0))
                    
                    # 写入文件
                    downloaded_size = 0
                    with open(file_path, 'wb') as f:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                                downloaded_size += len(chunk)
                                
                                # 显示下载进度
                                if total_size > 0 and downloaded_size % (1024 * 1024) == 0:
                                    progress = (downloaded_size / total_size) * 100
                                    logger.info(f"下载进度: {progress:.1f}% ({downloaded_size}/{total_size} bytes)")
                    
                    # 验证下载的文件
                    if file_path.exists() and file_path.stat().st_size > 0:
                        file_size = file_path.stat().st_size
                        logger.info(f"视频下载成功: {filename} (大小: {file_size} bytes)")
                        
                        result['success'] = True
                        result['local_path'] = str(file_path)
                        result['file_size'] = file_size
                        
                        # 记录到数据库
                        self._log_video_download(article_title, video_url, str(file_path), file_size, True)
                        
                        return result
                    else:
                        raise Exception("下载的文件大小为0")
                        
                except Exception as e:
                    logger.warning(f"下载尝试 {attempt + 1}/{self.max_retry} 失败: {e}")
                    if attempt < self.max_retry - 1:
                        time.sleep(5)  # 等待后重试
                        continue
                    else:
                        raise
            
        except Exception as e:
            error_msg = f"下载视频失败: {str(e)}"
            result['error'] = error_msg
            logger.error(error_msg)
            
            # 记录失败到数据库
            self._log_video_download(article_title, video_url, None, 0, False, error_msg)
        
        return result
    
    def _get_video_extension(self, url):
        """从URL中提取视频文件扩展名"""
        try:
            url_path = url.split('?')[0]  # 去除查询参数
            if '.' in url_path:
                ext = '.' + url_path.split('.')[-1].lower()
                valid_extensions = ['.mp4', '.m3u8', '.flv', '.avi', '.mov', '.mkv', '.webm']
                if ext in valid_extensions:
                    return ext
            return '.mp4'  # 默认扩展名
        except:
            return '.mp4'
    
    def _sanitize_filename(self, filename):
        """清理文件名，移除不安全字符"""
        if not filename:
            return "unnamed_article"
        
        # 移除或替换不安全字符
        unsafe_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*', '\n', '\r', '\t']
        safe_filename = filename
        for char in unsafe_chars:
            safe_filename = safe_filename.replace(char, '_')
        
        # 限制长度
        if len(safe_filename) > 100:
            safe_filename = safe_filename[:100]
            
        return safe_filename.strip()
    
    def _extract_title_from_url(self, url):
        """从URL中提取或生成文章标题"""
        try:
            # 尝试从URL路径中提取标题
            parsed_url = urlparse(url)
            if 's/' in parsed_url.path:
                # 微信文章URL通常包含文章ID
                article_id = parsed_url.path.split('s/')[-1]
                return f"article_{article_id[:10]}"
            else:
                return f"article_{int(time.time())}"
        except:
            return f"article_{int(time.time())}"
    
    def _log_video_download(self, article_title, video_url, local_path, file_size, success, error_message=None):
        """记录视频下载日志到数据库"""
        try:
            conn = pymysql.connect(**DB_CONFIG)
            with conn.cursor() as cursor:
                sql = """
                    INSERT INTO video_download_log 
                    (article_title, video_url, local_path, file_size, download_status, error_message, download_time)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                """
                status = 1 if success else 2  # 1=成功, 2=失败
                download_time = datetime.now() if success else None
                
                cursor.execute(sql, (
                    article_title, video_url, local_path, file_size, 
                    status, error_message, download_time
                ))
            conn.commit()
            conn.close()
            
        except Exception as e:
            logger.error(f"记录视频下载日志失败: {e}")
    
    def _log_download_summary(self, result):
        """记录下载摘要"""
        total_videos = len(result['video_urls'])
        successful_downloads = len(result['downloaded_videos'])
        failed_downloads = len(result['failed_videos'])
        
        logger.info("="*50)
        logger.info("视频下载摘要:")
        logger.info(f"文章标题: {result['article_title']}")
        logger.info(f"文章URL: {result['article_url']}")
        logger.info(f"发现视频数量: {total_videos}")
        logger.info(f"成功下载: {successful_downloads}")
        logger.info(f"下载失败: {failed_downloads}")
        
        if result['downloaded_videos']:
            logger.info("成功下载的视频:")
            for i, video in enumerate(result['downloaded_videos'], 1):
                logger.info(f"  {i}. {Path(video['local_path']).name} ({video['file_size']} bytes)")
        
        if result['failed_videos']:
            logger.info("下载失败的视频:")
            for i, video in enumerate(result['failed_videos'], 1):
                logger.info(f"  {i}. {video['url'][:50]}... (错误: {video['error']})")
        
        logger.info("="*50)





if __name__ == "__main__":
    url = "https://mp.weixin.qq.com/s/UkRQd4dVy7lmKyek_iwu8A"
    title = "测试文章"
    downloader = VideoDownloader()
    result = downloader.download_videos_from_article(url, title)